Add Distillation API to Keras #21572
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @divyashreepathihalli, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces a new Knowledge Distillation API to Keras, designed to facilitate the efficient transfer of learned knowledge from larger, pre-trained "teacher" models to smaller "student" models. The API seamlessly integrates with Keras's existing training, evaluation, and prediction workflows, providing a flexible and extensible framework for various distillation techniques.
Highlights
- New Distiller Model: A core Distiller class is added, which is a keras.Model subclass, enabling the combination and training of teacher and student models within the standard Keras workflow.
- Pluggable Distillation Strategies: Introduces a BaseDistillationStrategy and three concrete implementations: LogitsDistillation (for softening logits), FeatureDistillation (for intermediate feature matching), and MultiOutputDistillation (for handling models with multiple outputs).
- Configurable Loss Balancing: The Distiller allows specifying an alpha parameter to balance the contribution of the student's original loss and the distillation loss.
- Automatic Teacher Freezing: The teacher model is automatically set to non-trainable (trainable=False) during the distillation process to prevent its weights from being updated.
- Comprehensive Testing: New test files (distiller_test.py and strategies_test.py) are added to ensure the robustness and correctness of the new API, covering initialization, loss computation, and end-to-end workflows.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive Knowledge Distillation API to Keras, which is a great addition. The implementation is well-structured with a Distiller model and pluggable strategies. My review has identified a few issues: a critical issue with the FeatureDistillation strategy which is not fully implemented, a high-severity issue with an unused temperature parameter in the Distiller class that could mislead users, and a medium-severity issue regarding a simplistic fallback for loss calculation in multi-output scenarios. Addressing these points will improve the robustness and clarity of this new API.
keras/src/distillation/distiller.py
Outdated
| if isinstance(y_pred, (list, tuple)): | ||
| # For multi-output models, use the first output for student | ||
| # loss | ||
| # This is a simplified approach for compatibility | ||
| if isinstance(y, (list, tuple)): | ||
| student_loss = self.student_loss_fn(y[0], y_pred[0]) | ||
| else: | ||
| student_loss = self.student_loss_fn(y, y_pred[0]) |
There was a problem hiding this comment.
The fallback logic for calculating the student loss in _compute_loss for multi-output models is overly simplistic as it always defaults to using the first output (y_pred[0]). This might not align with user expectations for all multi-output scenarios and could lead to incorrect training behavior if model.compile() is not called with a loss that properly handles multiple outputs.
While the primary path using self.compiled_loss is correct, this fallback could be made more robust. Consider raising a more specific error if a multi-output model is used without a compiled loss, or clarifying this behavior more explicitly in the documentation.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a well-structured Knowledge Distillation API to Keras, a valuable addition for model compression and knowledge transfer. The implementation of the Distiller model and pluggable strategies is clean and the code is generally of high quality, with good documentation and examples.
I have identified a few significant issues that should be addressed to improve the robustness and completeness of the API:
- A critical flaw in
FeatureDistillationlimits its use to sequential models, which will prevent it from working with more complex architectures. - The
Distiller's loss computation for multi-output models has a bug in its fallback logic. - The
Distillermodel cannot be correctly serialized and deserialized due to an incompleteget_configimplementation. - The tests for
FeatureDistillationare not comprehensive enough to cover non-sequential model architectures.
Addressing these points will significantly enhance the reliability and usability of this new API. Overall, this is a great contribution.
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive Knowledge Distillation API to Keras, which is a great addition. The implementation is well-structured with a Distiller model and pluggable DistillationStrategy classes. The code is generally clean and includes a good set of tests.
I've identified a few key areas for improvement:
- Serialization: The
DistillerandMultiOutputDistillationclasses are not correctly configured for serialization, which will prevent models using them from being saved and loaded. This is a critical feature in Keras that needs to be addressed. - Feature Extraction Robustness: The
FeatureDistillationstrategy uses a method for extracting intermediate layer features that is not robust and will fail for models with non-sequential architectures (e.g., ResNets). This limitation needs to be documented and ideally improved. - Code Simplification: There's a small piece of unreachable code in the
Distiller's loss computation that can be simplified.
Addressing these points will significantly improve the robustness and usability of this new API.
keras/src/distillation/distiller.py
Outdated
| if isinstance(y_pred, list) and len(y_pred) > 0: | ||
| # For multi-output, use first output for student loss | ||
| student_loss = self.student_loss_fn(y[0], y_pred[0]) | ||
| else: | ||
| student_loss = self.student_loss_fn(y, y_pred) |
There was a problem hiding this comment.
The isinstance(y_pred, list) check on line 283 is redundant because y_pred is converted to a list on line 263. This makes the else block on line 286 unreachable. The logic can be simplified to directly use the first output for the student loss calculation.
# Fallback: use student_loss_fn directly
# For multi-output, use first output for student loss
student_loss = self.student_loss_fn(y[0], y_pred[0])
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive and well-designed Knowledge Distillation API to Keras. The implementation is robust, featuring a flexible Distiller class and a set of pluggable distillation strategies that cover common use cases like logits and feature distillation, as well as multi-output models. The code is accompanied by extensive and thorough tests, which is excellent. My feedback includes a couple of suggestions to improve code style in the API files and to enhance the robustness of a test case by removing a broad exception handler. Overall, this is a high-quality contribution that will be a valuable addition to Keras.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #21572 +/- ##
==========================================
- Coverage 82.59% 82.53% -0.06%
==========================================
Files 572 576 +4
Lines 58314 58846 +532
Branches 9130 9220 +90
==========================================
+ Hits 48166 48571 +405
- Misses 7817 7907 +90
- Partials 2331 2368 +37
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| if not isinstance(student_outputs, (list, tuple)): | ||
| student_outputs = [student_outputs] | ||
|
|
||
| if len(teacher_outputs) != len(student_outputs): |
There was a problem hiding this comment.
Oh, we agreed to support any structure for outputs, including dicts and nested outputs.
So remove lines 50-53 and replace lines 55-61 with keras.tree.assert_same_structures(teacher_outputs, student_outputs). It's actually a lot simpler.
There was a problem hiding this comment.
done!
| """ | ||
| if self._teacher_feature_extractor is not None: | ||
| # Use efficient multi-output extractor (returns dict directly) | ||
| return self._teacher_feature_extractor(x, training=False) |
There was a problem hiding this comment.
Hmm... so in the case when there are LogitsDistillation instances, that means that the student model is run twice. Once in self.call to get y_pred and once here to get the other logits.
There was a problem hiding this comment.
updated- the y_pred is now reused to avoid computation
There was a problem hiding this comment.
What I meant is that y_pred is computed once in self.call and then another time when calling _student_feature_extractor. But let's not worry about that now.
keras/src/distillation/distiller.py
Outdated
| except (ValueError, AttributeError): | ||
| # Fallback for subclassed models | ||
| return None |
There was a problem hiding this comment.
I think you should throw a meaningful error here.
If I read this correctly, when you try to use LogitsDistillation strategies on a subclass model (and all of KerasHub models are subclass models), this will return None, and then later, you'll get the error line 422 f"Layer '{layer_name}' features not found in extracted ", which will be confusing to users.
Instead you should say here (or somewhere else) that LogitsDistillation is not compatible with subclass models.
One alternative would be to check that in validate_model_compatibility. Right now, you validate that the layer name is found, but you could also validate that it's not a subclass model.
There was a problem hiding this comment.
KerasHub backbones are all written in functional model style. Logits distillation would work for subclass models but would not for feature distillation because model.inputs, model.outputs attribute would not be accessible.
Added checks in distiller and removed silent fallback
keras/src/distillation/distiller.py
Outdated
| f"targeting teacher layer " | ||
| f"'{strategy.teacher_layer_name}' and student " | ||
| f"layer '{strategy.student_layer_name}'. This can " | ||
| f"happen with subclassed models that haven't " |
There was a problem hiding this comment.
Doesn't it happen with all subclass models always? Even when they're built? (based on line 369)
Based on my comment line 369, can't we catch this error earlier rather than try to guess that any ValueError that's happening here has this as the cause?
keras/src/distillation/distiller.py
Outdated
| # Ensure distillation_loss is a scalar | ||
| if len(distillation_loss.shape) > 0: |
There was a problem hiding this comment.
Should this be an error (checking strategy_loss within the for loop)? Aren't strategies supposed to output scalars?
There was a problem hiding this comment.
raising an error inside the loop if strategy returns a non scalar
This PR adds Knowledge Distillation API to Keras,
Key Features
Core Components
Usage Examples
Basic Knowledge Distillation